package com.suxin.ratelimit.v2.handler;

import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.time.Duration;
import java.util.Collections;

import static com.suxin.ratelimit.v2.handler.RateLimitHandler.RateLimitConstKey.LIMIT_KEY;
import static com.suxin.ratelimit.v2.handler.RateLimitHandler.RateLimitConstKey.LUA_LIMIT;


/**
 * @author Tang
 * @classname DistributedRateLimitHandler
 * @description [  ]
 * @date 2022/2/10 13:48
 */
public class DistributedRateLimitHandler implements RateLimitHandler {

    @Resource
    private StringRedisTemplate redisTemplate;

    private DefaultRedisScript<Number> defaultRedisScript;

    @Override
    public boolean doHandler(Rate rate, String key) {
        Number number = this.redisTemplate.execute(defaultRedisScript,
                Collections.singletonList(String.join(":", LIMIT_KEY, key)),
                String.valueOf(rate.getLimit()), String.valueOf(Duration.of(rate.getDuration(), rate.getUnit()).getSeconds()));
        return number.intValue() == 0;
    }

    @PostConstruct
    public void init() {
        defaultRedisScript = new DefaultRedisScript<>();
        defaultRedisScript.setScriptText(LUA_LIMIT);
        defaultRedisScript.setResultType(Number.class);
    }

}